#ifndef Algebra_EigenSpace_h
#define Algebra_EigenSpace_h

#include "AlgebraModule.h"
#include "Space.h"

#include "Eigen/Dense"
#include "Eigen/Sparse"

SG_NAMESPACE_OPEN

namespace Algebra
{
    /**
     * @brief Wrap Eigen3 Vector.
     *
     * Convert the interface of a class into another interface clients expect.
     * Adapter lets classes work together that couldn't otherwise because of
     * incompatible interfaces.
     * A class adapter uses multiple inheritance to adapt one interface to another.
     *
     * @see GoF's Adpter Pattern.
     * @author nene
     * @date February, 2025.
     */
    template <typename TDataType>
    class ALGEBRA_EXPORT EigenVector : public Vector<TDataType>, private Eigen::VectorX<TDataType>
    {
      public:
        /// Default constructor
        EigenVector ();
        /// Construct a vector with given pattern.
        explicit EigenVector (Sparsity *sp);
        /// Construct a vector with given dimension.
        explicit EigenVector (const Index_t &n);
        /// Create holder of a PETSc Vec object/pointer. The Vec x object
        /// should already be created. If inc_ref_count is true, the reference
        /// counter of the Vec object will be increased. The Vec reference
        /// count will always be decreased upon destruction.
        EigenVector (Eigen::VectorX<TDataType> &x);
        /// Copy constructor
        EigenVector (const EigenVector<TDataType> &x);
        /// Move constructor
        EigenVector (EigenVector<TDataType> &&x);

        ~EigenVector () override;

        // Assignment operator
        EigenVector<TDataType> &operator= (const EigenVector<TDataType> &x);

        /** The corrresponding linear space which this algebraic object pertained to.*/
        Space<TDataType> *getSpace () override;

        /**
         * Calls the Vector's internal assembly routines, ensuring
         * that the values are consistent across processors.
         */
        void close () override;

        /**
         * Restores the \p Vector<T> to a pristine state.
         */
        void clear () override;

        /** Start assembling the elements stashed. */
        void beginAssembly () override;

        /** Finish assembling the elements stashed. */
        void endAssembly () override;

        /** Perform the assembly operation. */
        void assemble () override;

        // /** Test whether the vector has been assembled. */
        // bool assembled() override;

        /** Resize *this to a vector of length size. */
        void resize (Index_t size) override;

        /** Resize the vector to size while retaining old values. */
        void conservativeResize (Index_t size) override;

        /** The number of rows of the vector */
        Index_t rows () const override;

        /** The index of the first vector element stored on this processor.*/
        Index_t row_start () const override;

        /** The index of the last vector element (+1) stored on this processor. */
        Index_t row_end () const override;

        /** Set all entries to 0. */
        void setZero () override;

        /** Set v(i) = \p value. */
        void set (const Index_t &i, const TDataType &value) override;

        /** Add \p value to the vector entry specified by \p i. */
        void add (const Index_t &i, const TDataType &value) override;

        // /** The reference of the ith entry of the vector. */
        // virtual TDataType& operator() (const Index_t& i) = 0;

        /** A copy of the ith entry of the vector.*/
        TDataType operator() (const Index_t &i) const override;

        // /** The reference of the ith entry of the vector. */
        // virtual TDataType& operator[] (const Index_t& i) = 0;

        /** A copy of the ith entry of the vector.*/
        TDataType operator[] (const Index_t &i) const override;

        // /** The dot product of *this with other. */
        // TDataType dot(const PetscVector<TDataType> &) const override;

        /** Print the elements of a vector to the given output stream. */
        void print (std::ostream &out) const override;

        /** Create a subvector using the indices in \p rows. */
        void create_subvector (Vector<TDataType> &subvector, const GeneralIndexSet &rows) const override;

        /** Create a subvector from the stride index set in \p rows. */
        void create_subvector (Vector<TDataType> &subvector, const StrideIndexSet &rows) const override;

        /** Determine the vector component with maximum real part \p value and its location \p index. */
        void max (Index_t &index, TDataType &value) const override;

        /**
         * The \f$ \ell_{\infty} \f$-norm of the vector, i.e. the maximum
         * absolute value of the entries of the vector.
         */
        void linfty_norm (TDataType &value) const override;
    };

    /**
     * @brief Wrap Eigen3 Matrix.
     *
     * @author nene
     * @date February, 2025.
     */
    template <typename TDataType>
    class ALGEBRA_EXPORT EigenDenseMatrix : public DenseMatrix<TDataType>, private Eigen::Matrix<TDataType, Eigen::Dynamic, Eigen::Dynamic>
    {
    public:
        /** Print the elements of a matrix to the given output stream. */
        void print (std::ostream &out) const override;
    };

    /**
     * @brief Wrap Eigen3 SparseMatrix.
     *
     * @author nene
     * @date February, 2025.
     */
    template <typename TDataType>
    class ALGEBRA_EXPORT EigenSparseMatrix : public SparseMatrix<TDataType>
    {
      public:
        /// Default constructor
        EigenSparseMatrix ();
        /// Construct a matrix with given pattern.
        explicit EigenSparseMatrix (Sparsity *sp);
        EigenSparseMatrix (Eigen::SparseMatrix<TDataType> &A);
        /// Construct a vector with given dimension.
        EigenSparseMatrix (Index_t rows, Index_t cols);
        /// Copy constructor
        EigenSparseMatrix (const EigenSparseMatrix<TDataType> &A);
        /// Move constructor
        EigenSparseMatrix (EigenSparseMatrix<TDataType> &&A);
        /// Destructor
        ~EigenSparseMatrix () override;

        /// Assignment operator
        EigenSparseMatrix<TDataType> &operator= (const EigenSparseMatrix<TDataType> &A);

        /** The corrresponding linear space which this algebraic object pertained to.*/
        Space<TDataType> *getSpace () override;

        /** Restore the \p SparseMatrix<T> to a pristine state. */
        void clear () override;

        /**
         * Call the SparseMatrix's internal assembly routines, ensuring
         * that the values are consistent across processors.
         */
        void close () override;

        /** Start assembling the elements stashed. */
        void beginAssembly () override;

        /** Finish assembling the elements stashed. */
        void endAssembly () override;

        /** Perform the assembly operation. */
        void assemble () override;

        /** Test whether the matrix has been assembled. */
        bool assembled () const override;

        /** Resize the matrix to a rows x cols matrix. */
        void resize (Index_t rows, Index_t cols) override;

        /** Resize the matrix to a rows x cols matrix leaving old values untouched. */
        void conservativeResize (Index_t rows, Index_t cols) override;

        /** The global number of rows of the matrix */
        Index_t rows () const override;

        /** The global number of columns of the matrix */
        Index_t cols () const override;

        /** The index of the first matrix row stored on this processor.*/
        Index_t row_start () const override;

        /** The index of the last matrix row (+1) stored on this processor. */
        Index_t row_end () const override;

        /** The index of the first matrix column owned by this processor. */
        Index_t col_start () const override;

        /** The index of the last matrix column (+1) owned by this processor.*/
        Index_t col_end () const override;

        /** Set all entries to 0. */
        void setZero () override;

        /** Set *this to the identity matrix. */
        void setIdentity () override;

        /**
         * Set the element \p (i,j) to \p value.  Throws an error if the
         * entry does not exist. Zero values can be "stored" in non-existent
         * fields.
         */
        void set (const Index_t &i, const Index_t &j, const TDataType &value) override;

        /**
         * Add \p value to the element \p (i,j).  Throws an error if the
         * entry does not exist. Zero values can be "added" to non-existent
         * entries.
         */
        void add (const Index_t &i, const Index_t &j, const TDataType &value) override;

        // /**
        //  * Add the full matrix \p dm to the SparseMatrix.  This is useful
        //  * for adding an element matrix at assembly time.
        //  */
        // void add(const DenseMatrix<TDataType> &dm, const IndexMap_t &rows, const IndexMap_t &cols) override;

        // /**
        //  * Add \p spm to the rows and cols of this matrix (A)
        //  */
        // void add(const SparseMatrix<TDataType> & /*spm*/, const IndexMap_t & /*rows*/, const IndexMap_t & /*cols*/) override;

        // /** The reference of matrix entry \p(i,j). */
        // virtual TDataType& operator() (const Index_t& i, const Index_t& j) = 0;

        /** A copy of matrix entry \p (i,j). */
        TDataType operator() (const Index_t &i, const Index_t &j) const override;

        /** The number of nonzeros. */
        Index_t nonZeros () const override;

        /** Get the CSR (Compressed Sparse Row) Storage of the sparse matrix. */
        void get_csr (std::vector<Index_t> &ia, std::vector<Index_t> &ja, std::vector<TDataType> &a) const override;

        // /** Fill the matrix *this with the list of triplets defined by the iterator range begin - end. */
        // void setFromTriplets(const TripletIteratorType &begin, const TripletIteratorType &end);

        /** Print the elements of a matrix to the given output stream. */
        void print (std::ostream &out) const override;

        /** Create a submatrix from the row/col general-index-set. */
        void create_submatrix (SparseMatrix<TDataType> &submatrix, const GeneralIndexSet &rows, const GeneralIndexSet &cols) const override;

        /** Create a submatrix from the row/col stride-index-set. */
        void create_submatrix (SparseMatrix<TDataType> &submatrix, const StrideIndexSet &rows, const StrideIndexSet &cols) const override;

        /**
         * Similar to the above function, this function creates a \p
         * submatrix which is defined by the indices given in the \p rows
         * and \p cols vectors.
         * Note: Both \p rows and \p cols can be unsorted;
         *       Use the above function for better efficiency if your indices are sorted;
         *       \p rows and \p cols can contain indices that are owned by other processors.
         */
        void create_submatrix_nosort (SparseMatrix<TDataType> & /*submatrix*/,
                                      const GeneralIndexSet & /*rows*/,
                                      const GeneralIndexSet & /*cols*/) const override;

        /**
         * Create a submatrix from the row/col stride-index-set.
         * \p rows and \p cols can contain indices that are owned by other processors.
         *
         */
        void create_submatrix_nosort (SparseMatrix<TDataType> & /*submatrix*/,
                                      const StrideIndexSet & /*rows*/,
                                      const StrideIndexSet & /*cols*/) const override;

        /** Copy the diagonal part of the matrix into \p dest. */
        void get_diagonal (Vector<TDataType> &dest) const override;

        /** Copy the transpose of the matrix into \p dest, which may be *this. */
        void get_transpose (SparseMatrix<TDataType> &dest) const override;

        /** Copy the inverse of this matrix into \p dest, which may be *this. */
        void get_inverse (SparseMatrix<TDataType> &dest) const override;

        /** Scale all elements of this matrix by \p scale */
        void scale (const TDataType scale) override;

        /** Compute Y = A*X for matrix \p X. */
        void matrix_matrix_mult (SparseMatrix<TDataType> &X, SparseMatrix<TDataType> &Y) override;

      public:
        Eigen::SparseMatrix<TDataType, Eigen::RowMajor> mat;
    };

    template <typename TDataType>
    EigenVector<TDataType> operator+ (const EigenVector<TDataType> &, const EigenVector<TDataType> &);
    template <typename TDataType>
    EigenVector<TDataType> operator- (const EigenVector<TDataType> &, const EigenVector<TDataType> &);
    template <typename TDataType>
    EigenVector<TDataType> operator* (const TDataType &, const EigenVector<TDataType> &);
    template <typename TDataType>
    TDataType operator* (const EigenVector<TDataType> &, const EigenVector<TDataType> &);
    template <typename TDataType>
    EigenVector<TDataType> operator/ (const EigenVector<TDataType> &, const TDataType &);

    template <typename TDataType>
    EigenSparseMatrix<TDataType> operator+ (const EigenSparseMatrix<TDataType> &, const EigenSparseMatrix<TDataType> &);
    template <typename TDataType>
    EigenSparseMatrix<TDataType> operator- (const EigenSparseMatrix<TDataType> &, const EigenSparseMatrix<TDataType> &);
    template <typename TDataType>
    EigenSparseMatrix<TDataType> operator* (const TDataType &, const EigenSparseMatrix<TDataType> &);
    template <typename TDataType>
    EigenVector<TDataType> operator* (const EigenSparseMatrix<TDataType> &, const EigenVector<TDataType> &);

    /** Compute the matrix quadratic value \f $\boldsymbol{x}^{T}\boldsymbol{A}\boldsymbol{x} \f$ */
    template <typename TDataType>
    TDataType quadratic (const EigenSparseMatrix<TDataType> &A, const EigenVector<TDataType> &x);

    template <typename TDataType>
    std::ostream &operator<< (std::ostream &s, const EigenVector<TDataType> &v);
    template <typename TDataType>
    std::ostream &operator<< (std::ostream &s, const EigenDenseMatrix<TDataType> &v);
    template <typename TDataType>
    std::ostream &operator<< (std::ostream &s, const EigenSparseMatrix<TDataType> &m);

    /** Pack Eigen data types into a space. */
    template <typename TDataType>
    class ALGEBRA_EXPORT EigenSpace : public Space<TDataType>
    {
    public:
        using DataType         = TDataType;
        using VectorType       = Eigen::VectorX<TDataType>;
        using DenseMatrixType  = Eigen::MatrixX<TDataType>;
        using SparseMatrixType = Eigen::SparseMatrix<TDataType>;
        // using VectorType       = EigenVector<TDataType>;
        // using DenseMatrixType  = EigenDenseMatrix<TDataType>;
        // using SparseMatrixType = EigenSparseMatrix<TDataType>;
    };

    // Export the data types for external modules.
    template <typename TDataType>
    using Vector_t = typename EigenSpace<TDataType>::VectorType;
    template <typename TDataType>
    using Matrix_t = typename EigenSpace<TDataType>::DenseMatrixType;
    template <typename TDataType>
    using SparseMatrix_t = typename EigenSpace<TDataType>::SparseMatrixType;

    // real types
    typedef Vector_t<Real_t>       Vec_t;
    typedef Matrix_t<Real_t>       Mat_t;
    typedef SparseMatrix_t<Real_t> SpMat_t;

    // complex types
    typedef Vector_t<Complex_t>       ComplexVec_t;
    typedef Matrix_t<Complex_t>       ComplexMat_t;
    typedef SparseMatrix_t<Complex_t> ComplexSpMat_t;
  
  
    /**
     * @brief Replaces Eigen's setFromTriplets function template using MKL, needs to parse mat data
     * @tparam Scalar Matrix element type
     * @tparam TripletIteratorType Triplet iterator type
     * @param[in] begin Starting iterator of triplets
     * @param[in] end Ending iterator of triplets
     * @param[in out] mat Output sparse matrix (dimensions must be pre-set)
     * @note Requires linking with MKL library, calls mkl_sparse_?_create_coo interface
     * @author XuXihao
     * @date 2025.05.22
    */
    template <typename Scalar, typename TripletIteratorType>
    void convertTripletsToMKLSparse (const TripletIteratorType &begin, const TripletIteratorType &end, Eigen::SparseMatrix<Scalar> &mat);

    template <typename TDataType>
    bool isDiagonal (Eigen::SparseMatrix<TDataType> &mat, TDataType tolerance)
    {
        if (mat.rows () != mat.cols ())
            return false;
        for (int k = 0; k < mat.outerSize (); ++k)
        {
            for (Eigen::SparseMatrix<double>::InnerIterator it (mat, k); it; ++it)
            {
                if (it.row () != it.col () && fabs (it.value ()) > tolerance)
                {
                    return false;
                }
            }
        }
        return true;
    }
}  // namespace Algebra

SG_NAMESPACE_CLOSE

#endif  // Algebra_EigenSpace_h